// Copyright (c) 2020 Presto Labs Pte. Ltd.
// Author: donggu

#include <optional>
#include <stdexcept>

#include <LightGBM/c_api.h>
#include <torch/script.h>
#include <fmt/core.h>

#include "appcoin2/experimental/donggu/strat/feature_model.h"
#include "coin2/base/config/app_config.h"
#include "coin2/base/config/args.h"
#include "coin2/base/log.h"
#include "coin2/strategy/strategy.h"
#include "math.h"

using namespace coin2::exchange::base::symbology;
using namespace coin2::base::config;
using namespace coin2::exchange::feed;
using namespace coin2::app;

double thres = 0;
double lean_per_stack = 0;
int buy_y_idx = 0;
int sell_y_idx = 0;

struct Throttler {
  int64_t t = 0;
  const int64_t mn = 1;
  const int64_t mx = 1 * 1000000000LL;
  bool sample(int64_t ts) {
    if (t < ts) {
      int64_t delta = dist(mt) % (mx - mn + 1) + mn;
      t = ts + delta;
      return true;
    }
    return false;
  }
};

double get_available_qty(const IBook& book, double price, coin::proto::OrderDirection direction) {
  double sum = 0;
  if (direction == coin::proto::OrderDirection::DIRECTION_BUY) {
    for (const auto& [p, q] : book.Asks()) {
      if (p > price + 1e-8) {
        break;
      }
      sum += q;
    }
  } else {
    for (auto& [p, q] : book.Bids()) {
      if (p + 1e-8 < price) {
        break;
      }
      sum += q;
    }
  }
  return sum;
}

class LightModel {
 public:
  LightModel() {
    int out_num_iter;
    auto path = fmt::format("out/models/lgb_{}.txt", buy_y_idx);
    LGBM_BoosterCreateFromModelfile(
        path.c_str(),
        &out_num_iter,
        &buy_model_);
       path = fmt::format("out/models/lgb_{}.txt", sell_y_idx);
    LGBM_BoosterCreateFromModelfile(
        path.c_str(),
        &out_num_iter,
        &sell_model_);
  }
  std::pair<double, double> predict(double* features, int feature_size) {
    int64_t out_len;
    double out_results[2];
    LGBM_BoosterPredictForMatSingleRow(
        buy_model_,
        features,
        C_API_DTYPE_FLOAT64,
        feature_size,
        1,
        C_API_PREDICT_NORMAL,
        0,
        "",
        &out_len,
        out_results);
    if (out_len != 1) {
      throw std::runtime_error(std::string("lgbm predict() produced many values: ") + std::to_string(out_len));
    }
    LGBM_BoosterPredictForMatSingleRow(
        sell_model_,
        features,
        C_API_DTYPE_FLOAT64,
        feature_size,
        1,
        C_API_PREDICT_NORMAL,
        0,
        "",
        &out_len,
        &out_results[1]);
    if (out_len != 1) {
      throw std::runtime_error(std::string("lgbm predict() produced many values: ") + std::to_string(out_len));
    }
    return std::make_pair(out_results[0], out_results[1]);
  }
  ~LightModel() {
    LGBM_BoosterFree(buy_model_);
    LGBM_BoosterFree(sell_model_);
  }

  BoosterHandle buy_model_;
  BoosterHandle sell_model_;
  double out_buf_[2];
};

enum ModelType {
  Torch,
  Light,
};

class Strategy : public coin2::strategy::IStrategy {
 public:
  Strategy() {
    try {
      if (type_ == ModelType::Torch) {
        torch_model_ = torch::jit::load("out/models/traced.pt");
        torch_model_.eval();
        int found = 0;
        for (auto [name, buffer] : torch_model_.named_buffers()) {
          if (name == "y_mean") {
            found++;
            y_mean_ = buffer;
          } else if (name == "y_std") {
            found++;
            y_std_ = buffer;
          }
        }

        if (found != 2) {
          throw std::runtime_error("model does not have both buffers 'y_mean' and 'y_std'");
        }
      } else {
        light_model_.emplace();

      }

    } catch (const c10::Error& e) {
      throw std::runtime_error(std::string("error loading ML model: ") + e.what());
    }
  }

  void Init(::coin2::app::Driver* driver) override {
    multi_product_model_.Init(driver);
    target_product_holder_ = &product_cache_.holder(*multi_product_model_.target_);

    for (const auto& [name, oe] : driver->strategy()->order_executors()) {
      (void)name;
      oe_ = oe.get();
      break;
    }
    if (oe_ == nullptr) {
      throw std::runtime_error("no order executor found");
    }

    auto& systems = driver->order()->systems();
    if (systems.size() != 1) {
      throw std::runtime_error("driver config has more than 1 og");
    }
    oe_->set_order_system(systems[0].get());
  }

  void onTradeFeed(const FeedUpdate& upd) override {
    multi_product_model_.onTradeFeed(upd);
    onFeed(upd);
  }
  void onBookFeed(const FeedUpdate& upd) override {
    multi_product_model_.onBookFeed(upd);
    onFeed(upd);
  }

  void onFeed(const FeedUpdate& upd) {
    // any existing order to cancel?
    oe_->onPeriodicOrderCheck(*target_product_holder_, upd.timestamp());
    
    if (upd.is_book() && upd.product().key() == multi_product_model_.target_->key()) {
      target_book_ = &upd.book();
    }

    if (!multi_product_model_.IsReady(upd.timestamp())) {
      return;
    }

    double signal_ask_to_bid;
    double signal_bid_to_ask;

    if (type_ == ModelType::Torch) {
      // feed multi_product_model_.features_ into torch_model_
      torch::NoGradGuard _guard;

      // torch::Tensor x = torch::from_blob(multi_product_model_.features_,
      // {1,multi_product_model_.feature_size_}, torch::TensorOptions(torch::ScalarType::Float));
      torch::Tensor x = torch::empty({1, static_cast<long>(multi_product_model_.feature_size_)});
      for (size_t i = 0; i < multi_product_model_.feature_size_; i++) {
        x[0][i] = multi_product_model_.features_[i];
      }
      torch::Tensor y =
          torch_model_.forward({x}).toTensor() * y_std_.unsqueeze(0) + y_mean_.unsqueeze(0);

      double indicator = y[0][0].item().toDouble();
      signal_ask_to_bid = y[0][buy_y_idx].item().toDouble();
      signal_bid_to_ask = y[0][sell_y_idx].item().toDouble();

      if (indicator < 0.9 || indicator > 1.1) {  // overfit
        return;
      }

      // if (signal_ask_to_bid > 0.1) { return; }  // overfit
      // if (signal_bid_to_ask < -0.1) { return; }  // overfit
    } else {
      auto [buy, sell] = light_model_->predict(multi_product_model_.features_.get(), multi_product_model_.feature_size_);
      signal_ask_to_bid = buy;
      signal_bid_to_ask = sell;
      // printf("light signal %f %f\n", signal_ask_to_bid, signal_bid_to_ask);
    }

    /* --------------------------- print order status --------------------- */
    // static int64_t prev_submit_ts_ = 0;
    // if (prev_submit_ts_ + 20'000'000'000LL < upd.timestamp()) {
    //   LOG(INFO) << multi_product_model_.target_->product();
    //   LOG(INFO) << oe_->GetLiveOrderSummaryFormattedString(upd.timestamp());
    //   prev_submit_ts_ = upd.timestamp();
    // }

    /* ----------------------------- pick up signal ----------------------------- */
    auto ask0p = multi_product_model_.target_model_->features_[Ask0p];
    auto bid0p = multi_product_model_.target_model_->features_[Bid0p];
    double sign = 0;
    double rounded_price = NAN;
    OrderDirection direction;

    // double current_stack = oe_->GetCurrentStack(product_holder); // TODO

    // calc lean threshold
    double current_stack = oe_->GetCurrentStack(*target_product_holder_);
    // the sign is different from agg logic
    double buy_thres = thres + (current_stack > 0) * current_stack * lean_per_stack;
    double sell_thres = -thres + (current_stack < 0) * current_stack * lean_per_stack;

    if (signal_ask_to_bid > buy_thres) {  // AskToBid > thres
      sign += 1;
      direction = OrderDirection::DIRECTION_BUY;
      rounded_price = ask0p;
    }
    if (signal_bid_to_ask < -sell_thres) {  // BidToAsk < -thres
      sign -= 1;
      direction = OrderDirection::DIRECTION_SELL;
      rounded_price = bid0p;
    }

    if (sign == 0) {
      // both buy&sell signals are strong?
      // skip
      return;
    }

    // printf("signal %10.3f and %10.3f\n",signal_ask_to_bid*10000.0, signal_bid_to_ask*10000.0);

    if (!throttle_.sample(upd.timestamp())) {
      return;
    }

    /* --------------------------------- submit --------------------------------- */

    double raw_qty = oe_->GetConfig(*target_product_holder_).qty().max_lot_size();
    raw_qty = std::min(raw_qty, get_available_qty(*target_book_, rounded_price, direction));
    const double leverage_rate = 10.0;  // source: jshin
    const bool post_only = false;

    if (raw_qty > 0) {
      oe_->SubmitOrder(
          *target_product_holder_,
          rounded_price,  // OE does rounding, so this can be replaced by
                          // ExecutorConfig-RoundingPolicy
          raw_qty,
          direction,
          OrderDuration::GTC_ORDER,
          OrderType::LIMIT_ORDER,
          std::nullopt,
          leverage_rate,
          post_only,
          "",
          upd.timestamp());

      nlohmann::json msg = {
          {"compartment", "strategy"},
          {"component", "donggu-exp-logic"},
          {"product", multi_product_model_.target_->product()},
          {"order_sign", sign},
          {"order_raw_qty", raw_qty},
          {"price_order", rounded_price}};

      LOG(INFO) << msg.dump();
    }
  }

 private:
  torch::jit::script::Module torch_model_;
  torch::Tensor y_mean_;
  torch::Tensor y_std_;

  ModelType type_ = ModelType::Light;

  std::optional<LightModel> light_model_;
  MultiProductModel multi_product_model_;
  const IProductHolder* target_product_holder_ = nullptr;

  Throttler throttle_;
  ::coin2::strategy::order_executor::OrderExecutorSystem* oe_ = nullptr;
  const IBook* target_book_ = nullptr;
  ProductEncyclopedia product_cache_;
};

int main(int argc, char* argv[]) {
  DriverConfig app;

  cxxopts::Options opt("experimental strat", "experimental strat");
  opt.add_options()("thres", "threshold", cxxopts::value<double>());
  opt.add_options()("horizon", "horizon", cxxopts::value<int>());
  opt.add_options()("sim_result_filepath", "sim_result_filepath", cxxopts::value<std::string>());


  AddDriverOptions(&opt);
  AddIntervalOptions(&opt);

  auto res = opt.parse(argc, argv);
  if (res.count("thres") == 0) {
    throw std::runtime_error("provide --thres");
  }
  if (res.count("horizon") == 0) {
    throw std::runtime_error("provide --horizon");
  }
  if (res.count("lean_per_stack") == 0) {
    throw std::runtime_error("provide --lean_per_stack");
  }
  thres = res["thres"].as<double>();
  lean_per_stack = res["lean_per_stack"].as<double>();
  int horizon = res["horizon"].as<int>();
  if (horizon == 3) {
    buy_y_idx = 1;
  } else if (horizon == 5) {
    buy_y_idx = 2;
  } else if (horizon == 10) {
    buy_y_idx = 3;
  } else if (horizon == 30) {
    buy_y_idx = 4;
  } else {
    throw std::runtime_error("horizon must be one of [3, 5, 10, 30]");
  }
  sell_y_idx = buy_y_idx + 4;

  ParseDriverOptions(res, &app);
  ParseIntervalOptions(
      res,
      app.mutable_feed()->mutable_common()->mutable_archive()->mutable_interval());
  auto interval = FromIntervalConfig(app.feed().common().archive().interval());
  app.mutable_order()->mutable_common()->mutable_sim_config()->set_start_timestamp(interval.first);
  ParseProductOptions(res, app.mutable_feed());

  if (res.count("sim_result_filepath") > 0) {
    auto* mut_sim_config = app.mutable_order()->mutable_common()->mutable_sim_config();
    mut_sim_config->set_sim_result_filepath(res["sim_result_filepath"].as<std::string>());
  }

  Strategy s;  // simulation
  Driver driver(app, &s);
  driver.Run();

  return 0;
}
